/*
 * Copyright (c) 2025, 大连理工大学 (Dalian University of Technology)
 *
 * Licensed under the Mulan PSL v2.
 * You can use this file according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of the License at
 *
 *     http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE. SEE THE MULAN PSL v2
 * FOR MORE DETAILS.
 */

#pragma once

#include "DBManager/IRepository/Base/IMainlyRepository.h"
#include "DBManager/Repository/HDF5/Base/CacheRepositoryBase.h"

SG_DBMANAGER_HDF5_NAMESPACE_OPEN

template <typename T>
class DBManager_DECLSPEC MainlyRepositoryBase : public CacheRepositoryBase<T>, public virtual IMainlyRepository<T>
{
  public:
    explicit MainlyRepositoryBase (HDF5Database& database) : IMainlyRepository<T> (), CacheRepositoryBase<T> (database) {};

    ~MainlyRepositoryBase () = default;

    using CacheRepositoryBase<T>::Sort;
    using CacheRepositoryBase<T>::GetDatabase;
    using CacheRepositoryBase<T>::GetDataset;
    using CacheRepositoryBase<T>::GetDatasetInfo;
    using CacheRepositoryBase<T>::GetMainDatasetInfo;
    using CacheRepositoryBase<T>::GetIndexDatasetInfo;
    using CacheRepositoryBase<T>::DATASET;
    using CacheRepositoryBase<T>::GetIndexCache;

    /**
     * @brief    Get the Current Area IDs object
     *
     * @param    ids
     * @return
     * @author   wanglong
     * @date     2025.08.05
     */
    virtual void IdFilter (std::vector<SG::DataStructure::Common::Id>& ids)
    {}

    /**
     * @brief 获取表单的行数（数据总量）。
     *
     * @return 表单的行数。
     */
    virtual std::size_t Size () override
    {
        Sort ();

        if (this->GetMainDatasetInfo ().IsCreate ())
        {
            return GetDataset (GetIndexDatasetInfo ()).getElementCount ();  ///< 读取所有 ID
        }
        return 0;
    }

    /**
     * @brief 向缓存中添加多个实体。当缓存达到最大数量时，将调用 `Dump()` 方法进行处理。
     *
     * @param entity 要添加的实体。
     */
    SG::DataStructure::Common::Status Add (std::vector<T>& entity) override
    {
        auto& mainCache = this->GetMainCache ();

        mainCache.m_writeStorage.insert (mainCache.m_writeStorage.end (), entity.begin (), entity.end ());

        // 如果缓存已满，则执行 Dump 操作
        if (mainCache.m_writeStorage.size () >= mainCache.m_limitSize)
        {
            this->DumpWriteCahce ();
        }
        GetMainDatasetInfo ().SetIsSequential (false);

        return SG::DataStructure::Common::Status::STATUS_Success;
    }

    /**
     * @brief 向缓存中添加一个实体。当缓存达到最大数量时，将调用 `Dump()` 方法进行处理。
     *
     * @param entity 要添加的实体。
     */
    SG::DataStructure::Common::Status Add (T& entity) override
    {
        auto& mainCache = this->GetMainCache ();

        mainCache.m_writeStorage.push_back (entity);  ///< 将实体添加到写缓存中

        // 如果缓存已满，则执行 Dump 操作
        if (mainCache.m_writeStorage.size () >= mainCache.m_limitSize)
        {
            this->DumpWriteCahce ();
        }

        // 数据添加后对数据的处理
        this->AfterAdd (entity);

        GetMainDatasetInfo ().SetIsSequential (false);

        return SG::DataStructure::Common::Status::STATUS_Success;
    }

    /**
     * @brief 获取所有的 ID。
     *
     * @return 包含所有 ID 的向量。
     */
    std::vector<SG::DataStructure::Common::Id> GetAllId () override
    {
        Sort ();

        std::vector<SG::DataStructure::Common::Id> idList;
        if (this->GetMainDatasetInfo ().IsCreate ())
        {
            GetDataset (GetIndexDatasetInfo ()).read (idList);  ///< 读取所有 ID
        }
#ifdef SG_USE_MPI
        IdFilter (idList);
#endif

        return idList;
    }

    /**
     * @brief 计算指定 ID 的出现次数。
     *
     * @param id 要计算的 ID。
     * @return ID 出现的次数。
     */
    std::size_t Count (SG::DataStructure::Common::Id id) override
    {
        Sort ();

        auto ids = GetAllId ();  ///< 获取所有 ID
        auto it  = std::lower_bound (ids.begin (), ids.end (), id);
        if (it == ids.end () || *it != id)
            return false;
        return true;
    }

    /**
     * @brief 根据索引查找实体。
     *
     * @param index 要查找的索引(必须确定是全局索引)。
     * @param _OUT entity 查找到的实体将存储在此参数中。
     * @return 如果索引有效且实体成功加载，则返回 `true`；否则返回 `false`。
     *
     * @throw std::runtime_error 如果索引超出范围。
     */
    virtual SG::DataStructure::Common::Status FindByIndex (::std::size_t index, T& _OUT entity)
    {
        Sort ();

        // 检查索引是否超出范围
        if (GetMainDatasetInfo ().GetRowCount () <= index)
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        auto& mainCache = this->GetMainCache ();

        std::vector<T>* cacheVec;

        if (index >= mainCache.dataMaxIndex || index < mainCache.dataMinIndex
            || (mainCache.dataMinIndex == 0 && mainCache.dataMaxIndex == 0))
        {
            cacheVec = this->LoadReadCache (index);
        }
        else
        {
            cacheVec = this->m_cacheHelper.template GetVectorFast<T> (mainCache.m_readCacheKey);
            // 说明缓存失效了，需要重新读取到缓存中
            if (cacheVec == nullptr)
            {
                cacheVec = this->LoadReadCache (index);
            }
        }

        if (cacheVec == nullptr)
        {
            SGError << GetMainDatasetInfo ().GetPath () << " Failed to load read cache for index: " << index;
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }
        else
        {
            entity = (*cacheVec)[index - mainCache.dataMinIndex];  ///< 获取实体
            return SG::DataStructure::Common::Status::STATUS_Success;
        }
    }

    /**
     * @brief
     *
     * @param id 要查找的 ID。
     * @param _OUT entity 查找到的实体将存储在此参数中。
     * @return 如果 ID 成功找到并加载了实体，则返回 `true`；否则返回 `false`。
     *
     * @throw std::runtime_error 如果 ID 不存在或输入无效。
     */
    virtual SG::DataStructure::Common::Status FindById (SG::DataStructure::Common::Id id, T& _OUT entity) override
    {
        // 1. 确定数据是排序过的
        Sort ();

        if (!this->GetMainDatasetInfo ().IsCreate ())
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        // 2. 在缓存中查找对应ID的数据 使用二分查找
        // tip. 测试证明std::lower_bound 在 100w与10w数据量情况下 使用vector<int>会比vector<Struct>快6%左右。
        auto& indexCache = GetIndexCache ();

        std::unordered_map<SG::DataStructure::Common::Id, std::size_t>* cacheStorage =
            this->m_cacheHelper.template GetUnorderedMapFast<SG::DataStructure::Common::Id, std::size_t> (indexCache.cacheKey);

        if (cacheStorage != nullptr)
        {
            auto iter = cacheStorage->find (id);
            // 2.1 如果找到了根据索引取对应的数据
            if (iter != cacheStorage->end ())
            {
                return MainlyRepositoryBase::FindByIndex ((iter->second + indexCache.minIndex), entity);
            }
        }

        // 2.2 如果没有找到，去文件中查找
        auto& indexTableInfo = GetIndexDatasetInfo ();

        // 2.2.1. 读取所有 ID
        std::vector<SG::DataStructure::Common::Id> idList;
        GetDatabase ().GetFile ()->getDataSet (indexTableInfo.GetPath ()).read (idList);

        // 2.2.2. 在所有ID中进行二分查找
        auto it = std::lower_bound (idList.begin (), idList.end (), id);
        //    所有ID中都没有找到
        if (it == idList.end () || *it != id)
        {
            return SG::DataStructure::Common::Status::STATUS_Failure;
        }

        // 2.2.3. 计算 ID 在表单中的索引
        const size_t index = std::distance (idList.begin (), it);

        // 2.2.5. 更新 indexCahce
        std::size_t beginIndex = 0, length = 0;
        indexTableInfo.CalcIndexBeginAndLength (index, _OUT beginIndex, _OUT length, this->GetDatasetCacheSize ());  ///< 获取缓存范围

        cacheStorage = new std::unordered_map<SG::DataStructure::Common::Id, std::size_t> ();

        for (std::size_t i = 0; i < length; ++i)
        {
            (*cacheStorage)[idList[beginIndex + i]] = beginIndex + i;
        }
        this->m_cacheHelper.template CacheUnorderedMap<SG::DataStructure::Common::Id, std::size_t> (indexCache.cacheKey, cacheStorage);

        indexCache.minIndex = beginIndex;
        indexCache.maxIndex = indexCache.minIndex + length;

        // 2.2.6. 返回数据
        return MainlyRepositoryBase::FindByIndex (index, entity);
    }

#if defined(SG_USE_MPI)
    /**
     * @brief    过滤函数，获取当前分区内的id
     *
     * @param    ids
     * @return
     * @author   wanglong
     * @date     2025.08.06
     */
    static void ElementIdFilter (IDatabase& database, std::vector<SG::DataStructure::Common::Id>& ids)
    {
        if (database.GetPartitionisFilter ())
        {
            const std::unordered_set<SG::DataStructure::Common::Id>& areaIdList = database.GetPartitionElement ();  // 获取当前进程的ID列表
            std::vector<SG::DataStructure::Common::Id>               temp       = std::move (ids);
            ids.clear ();                                                                                           // 清空原有的idList
            ids.reserve (areaIdList.size ());                                                                       // 预分配空间以提高性能
            for (auto&& id : temp)
            {
                if (areaIdList.find (id) != areaIdList.end ())
                {
                    ids.emplace_back (id);  // 仅添加当前进程的ID
                }
            }
        }
    }
#endif
};

SG_DBMANAGER_HDF5_NAMESPACE_CLOSE
